{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "import subprocess\n",
    "\n",
    "from openai import OpenAI\n",
    "import sys\n",
    "sys.path.append('../')\n",
    "from prompts.design_prompt import design_prompt\n",
    "from prompts.parsing_prompt import parsing_prompt\n",
    "from prompts.scripting_prompt import scripting_prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../config/config.json') as f:\n",
    "     config = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "OPENSCAD_EXEC = config[\"OPENSCAD_EXEC_PATH\"]\n",
    "FILE_DIR = \"../common/spur_gear\"\n",
    "DESIGN_FILENAME = \"gears_design.md\"\n",
    "PARAMETERS_FILENAME = \"gears_parameters.json\"\n",
    "SCRIPT_FILENAME = \"gears_script.scad\"\n",
    "if not os.path.exists(FILE_DIR):\n",
    "    os.makedirs(FILE_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = OpenAI(api_key=config['OPENAI_API_KEY'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "design_description = f\"\"\" \n",
    "Suppose there is a one-stage system with two spur gears.\n",
    "Gear axes are parallel and the distance between two axes is 100mm. \n",
    "The height of driven and driving gears are 27mm and 18mm respectively, the pressure angle are both 20 degrees.\n",
    "The desired gear ratio of the reduction system is 3:1.\n",
    "Set the driving gear to the center of the plane, which is the coordinate (0, 0)\n",
    "Based on the given information, computer module number, teeth number of both gears and the coordniate of the driven gear centers.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "design_messages=[\n",
    "        {\"role\": \"system\", \"content\": design_prompt},\n",
    "        {\"role\": \"user\", \"content\": design_description},\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# stop_tokens = ['']\n",
    "# collected_tokens = []\n",
    "\n",
    "# res = client.chat.completions.create(\n",
    "#     model=config['GPT_MODEL'],\n",
    "#     messages=messages,\n",
    "#     stream=True,\n",
    "# )\n",
    "\n",
    "# for chunk in res:\n",
    "#     temp_tokens = chunk.choices[0].delta.content\n",
    "#     if temp_tokens not in stop_tokens:\n",
    "#         collected_tokens.append(temp_tokens)\n",
    "#     else:\n",
    "#         chunk_res = \"\".join(collected_tokens).strip()\n",
    "#         print(chunk_res)\n",
    "#         collected_tokens = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = client.chat.completions.create(\n",
    "    model=config['GPT_MODEL'],\n",
    "    messages=design_messages,\n",
    "    temperature=0.1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "To design the gear system, we need to determine the module, number of teeth for both gears, and the coordinates of the driven gear center. Let's start by calculating the module and number of teeth for both gears to achieve the desired gear ratio of 3:1.\n",
      "\n",
      "### Step 1: Gear Ratio and Number of Teeth\n",
      "The gear ratio (i) is defined as the ratio of the number of teeth on the driven gear (N2) to the number of teeth on the driving gear (N1):\n",
      "\n",
      "$$ i = \\frac{N2}{N1} $$\n",
      "\n",
      "Given the desired gear ratio is 3:1, we can express N2 as 3 times N1:\n",
      "\n",
      "$$ N2 = 3N1 $$\n",
      "\n",
      "### Step 2: Module Calculation\n",
      "The module (m) is a measure of the size of the teeth and is related to the pitch diameter (d) and the number of teeth (N) by the formula:\n",
      "\n",
      "$$ m = \\frac{d}{N} $$\n",
      "\n",
      "The center distance (a) between the two gears can be expressed in terms of the module and the number of teeth:\n",
      "\n",
      "$$ a = \\frac{m(N1 + N2)}{2} $$\n",
      "\n",
      "Given that the distance between the gear axes is 100mm, we can set this equal to the center distance:\n",
      "\n",
      "$$ 100mm = \\frac{m(N1 + 3N1)}{2} $$\n",
      "$$ 100mm = 2mN1 $$\n",
      "\n",
      "We can choose an initial value for N1 and calculate the module (m). Let's choose a standard value for N1, such as 20 teeth for the driving gear:\n",
      "\n",
      "$$ m = \\frac{100mm}{2 \\times 20} $$\n",
      "$$ m = \\frac{100mm}{40} $$\n",
      "$$ m = 2.5mm $$\n",
      "\n",
      "### Step 3: Check Module Availability\n",
      "The calculated module is 2.5mm, which is a standard value. Therefore, we can use this module for our gears.\n",
      "\n",
      "### Step 4: Calculate Number of Teeth for Driven Gear\n",
      "Now that we have the module and the number of teeth for the driving gear, we can calculate the number of teeth for the driven gear:\n",
      "\n",
      "$$ N2 = 3N1 $$\n",
      "$$ N2 = 3 \\times 20 $$\n",
      "$$ N2 = 60 $$\n",
      "\n",
      "### Step 5: Calculate Pitch Diameters\n",
      "With the module and number of teeth, we can calculate the pitch diameters for both gears:\n",
      "\n",
      "$$ d1 = mN1 $$\n",
      "$$ d1 = 2.5mm \\times 20 $$\n",
      "$$ d1 = 50mm $$\n",
      "\n",
      "$$ d2 = mN2 $$\n",
      "$$ d2 = 2.5mm \\times 60 $$\n",
      "$$ d2 = 150mm $$\n",
      "\n",
      "### Step 6: Determine Coordinates of Driven Gear\n",
      "The driving gear is at the center of the plane (0,0). The driven gear needs to be positioned such that its center is 100mm away from the driving gear. Since the gears are on the same horizontal line, the y-coordinate will be the same (0). The x-coordinate will be the sum of the radii of both gears:\n",
      "\n",
      "$$ x = \\frac{d1}{2} + \\frac{d2}{2} $$\n",
      "$$ x = \\frac{50mm}{2} + \\frac{150mm}{2} $$\n",
      "$$ x = 25mm + 75mm $$\n",
      "$$ x = 100mm $$\n",
      "\n",
      "Therefore, the coordinate of the driven gear center is (100, 0).\n",
      "\n",
      "### Final Result:\n",
      "gear 1 (driving gear): spur gear, module 2.5mm, teeth 20, pitch diameter 50mm, height 18mm, coordinate (0,0)\n",
      "gear 2 (driven gear): spur gear, module 2.5mm, teeth 60, pitch diameter 150mm, height 27mm, coordinate (100,0)\n"
     ]
    }
   ],
   "source": [
    "print(res.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(FILE_DIR, DESIGN_FILENAME), \"w\") as f:\n",
    "    f.write(res.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "parsing_description = \\\n",
    "f\"gear 1 (driving gear): spur gear, module 2.5mm, teeth 20, pitch diameter 50mm, height 18mm, coordinate (0,0)\" + \\\n",
    "f\"pressure angle 20 degrees \\n\" + \\\n",
    "f\"gear 2 (driven gear): spur gear, module 2.5mm, teeth 60, pitch diameter 150mm, height 27mm, coordinate (100,0)\" + \\\n",
    "f\"pressure angle 20 degrees\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "parsing_messages = [\n",
    "    {\"role\": \"system\", \"content\": parsing_prompt},\n",
    "    {\"role\": \"user\", \"content\": parsing_description},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "parsing_response = client.chat.completions.create(\n",
    "    model=config['GPT_MODEL'],\n",
    "    response_format={ \"type\": \"json_object\" },\n",
    "    messages=parsing_messages,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"gear_1\": {\n",
      "    \"gear_type\": \"spur\",\n",
      "    \"source\": \"driving\",\n",
      "    \"coordinate\": {\n",
      "      \"x\": 0,\n",
      "      \"y\": 0,\n",
      "      \"z\": 0,\n",
      "      \"alpha\": 0,\n",
      "      \"beta\": 0,\n",
      "      \"gamma\": 0\n",
      "    },\n",
      "    \"unit\": \"mm\",\n",
      "    \"module\": 2.5,\n",
      "    \"teeth\": 20,\n",
      "    \"height\": 18,\n",
      "    \"pitch_d\": 50,\n",
      "    \"helix_angle\": 0,\n",
      "    \"pitch_angle\": 0,\n",
      "    \"pressure_angle\": 20\n",
      "  },\n",
      "  \"gear_2\": {\n",
      "    \"gear_type\": \"spur\",\n",
      "    \"source\": \"driven\",\n",
      "    \"coordinate\": {\n",
      "      \"x\": 100,\n",
      "      \"y\": 0,\n",
      "      \"z\": 0,\n",
      "      \"alpha\": 0,\n",
      "      \"beta\": 0,\n",
      "      \"gamma\": 0\n",
      "    },\n",
      "    \"unit\": \"mm\",\n",
      "    \"module\": 2.5,\n",
      "    \"teeth\": 60,\n",
      "    \"height\": 27,\n",
      "    \"pitch_d\": 150,\n",
      "    \"helix_angle\": 0,\n",
      "    \"pitch_angle\": 0,\n",
      "    \"pressure_angle\": 20\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "gears_parameters = parsing_response.choices[0].message.content.strip()\n",
    "print(gears_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(FILE_DIR, PARAMETERS_FILENAME), \"w\") as f:\n",
    "    f.write(gears_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(FILE_DIR, PARAMETERS_FILENAME), \"r\") as f:\n",
    "    parameters = json.load(f)\n",
    "parameters = str(parameters)\n",
    "\n",
    "scripting_description = f\"Write a script according to provided parameters of gears \\n\" + parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "script_messages = [\n",
    "    {\"role\": \"system\", \"content\": scripting_prompt},\n",
    "    {\"role\": \"user\", \"content\": scripting_description}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"../common/spur_gear/gears_parameter.json\", \"r\") as f:\n",
    "    parameters = json.load(f)\n",
    "parameters = str(parameters)\n",
    "\n",
    "scripting_description = f\"Write a script according to provided procedure with parameters of gears\" + parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Write a script according to provided parameters of gears \n",
      "{'gear_1': {'gear_type': 'spur', 'source': 'driving', 'coordinate': {'x': 0, 'y': 0, 'z': 0, 'alpha': 0, 'beta': 0, 'gamma': 0}, 'unit': 'mm', 'module': 2.5, 'teeth': 20, 'height': 18, 'pitch_d': 50, 'helix_angle': 0, 'pitch_angle': 0, 'pressure_angle': 20}, 'gear_2': {'gear_type': 'spur', 'source': 'driven', 'coordinate': {'x': 100, 'y': 0, 'z': 0, 'alpha': 0, 'beta': 0, 'gamma': 0}, 'unit': 'mm', 'module': 2.5, 'teeth': 60, 'height': 27, 'pitch_d': 150, 'helix_angle': 0, 'pitch_angle': 0, 'pressure_angle': 20}}\n"
     ]
    }
   ],
   "source": [
    "print(scripting_description)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "script_messages = [\n",
    "    {\"role\": \"system\", \"content\": scripting_prompt},\n",
    "    {\"role\": \"user\", \"content\": scripting_description}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "scripting_response = client.chat.completions.create(\n",
    "    model=config['GPT_MODEL'],\n",
    "    messages=script_messages,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def script_postprocessing(script):\n",
    "    # To remove textual perfix or suffix information in the raw script generated by GPT-4\n",
    "    pattern = re.compile(r'```(?:openscad|scad)(.*?)```', re.DOTALL)\n",
    "    matches = re.findall(pattern, script)\n",
    "    clean_script = \"\"\n",
    "    if len(matches) > 0:\n",
    "        for m in matches:\n",
    "            clean_script = clean_script + m.strip()\n",
    "    else:\n",
    "        clean_script = script\n",
    "    return clean_script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "translate([0,0,0]) rotate([0,0,0])\n",
      "gear(m=2.5, z=20, h=18, w=20);\n",
      "\n",
      "translate([100,0,0]) rotate([0,0,(360/(60*2))])\n",
      "gear(m=2.5, z=60, h=27, w=20);\n"
     ]
    }
   ],
   "source": [
    "script = scripting_response.choices[0].message.content.strip()\n",
    "script = script_postprocessing(script)\n",
    "print(script)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(FILE_DIR, SCRIPT_FILENAME), \"w\") as f:\n",
    "    f.write(\"include <gears.scad>\")\n",
    "\n",
    "with open(os.path.join(FILE_DIR, SCRIPT_FILENAME), \"a\") as f:\n",
    "    f.write(\"\\n\" + script)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Popen: returncode: None args: ['D:\\\\openSCAD\\\\install\\\\openscad.exe', '../c...>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "subprocess.Popen([OPENSCAD_EXEC, os.path.join(FILE_DIR, SCRIPT_FILENAME)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chatbot",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
